package com.cneport.das.common.security.pub;

import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
import sun.security.pkcs.PKCS7;

import java.io.InputStreamReader;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPublicKeySpec;

/**
 * 解析及生成密钥工具类
 */
public class DefaultKeyStore {
    PrivateKey privateKey;
    PublicKey publicKey;
    java.security.cert.Certificate cert, certChain[];

    public DefaultKeyStore() {
    }

    public DefaultKeyStore(PrivateKey privateKey, PublicKey publicKey) {
        this.privateKey = privateKey;
        this.publicKey = publicKey;
    }

    protected static String getKeyStoreTypeByFileSuffix(String url) {

        if (url != null) {
            int p = url.lastIndexOf('/');
            String fileName = p < 0 ? url : url.substring(p + 1);
            p = fileName.lastIndexOf('.');
            String suffix = p < 0 ? null : fileName.substring(p + 1);
            if (suffix != null) switch (suffix.toLowerCase()) {
                case "pfx":
                case "p12":
                    return "PKCS12";
                case "cer":
                    return "X.509";
                case "jks":
                    return "JKS";
                case "pem":
                    return "PEM";
            }

        }
        return "JKS";
    }


    public void load(String storeUrl, String storeType) {
        load(storeUrl, storeType, null, null, null);
    }


    /**
     * @param storeUrl
     * @param storeType
     * @param storePass
     * @param keyAlias
     * @param keyPass
     */
    public void load(String storeUrl, String storeType,
                     final String storePass,
                     final String keyAlias, final String keyPass) {
        try {
            load(new java.net.URL(storeUrl), storeType, storePass, keyAlias, keyPass);
        } catch (java.net.MalformedURLException ex) {
            throw new RuntimeException(ex);
        }
    }

    public void load(java.net.URL storeUrl, String storeType,
                     final String storePass,
                     final String keyAlias, final String keyPass) {
        if (storeType == null)
            storeType = getKeyStoreTypeByFileSuffix(storeUrl.toString());
        try (
                java.io.InputStream fis = storeUrl.openStream();
        ) {
            switch (storeType) {
                case "X.509": {
                    this.cert = loadX509Certificate(fis);
                    this.publicKey = this.cert.getPublicKey();
                }
                return;

                case "PEM": {
                    try (
                            PEMParser PEMParser = new PEMParser(new InputStreamReader(fis));
                    ) {
                        Object pemObj = PEMParser.readObject();
                        if (pemObj instanceof PEMKeyPair)
                            pemObj = new JcaPEMKeyConverter().getKeyPair((PEMKeyPair) pemObj);
                        if (pemObj instanceof KeyPair) {
                            this.privateKey = ((KeyPair) pemObj).getPrivate();
                            this.publicKey = ((KeyPair) pemObj).getPublic();
                        } else if (pemObj instanceof SubjectPublicKeyInfo) {
                            this.publicKey = new JcaPEMKeyConverter().getPublicKey((SubjectPublicKeyInfo) pemObj);
                        } else if (pemObj instanceof PrivateKeyInfo) {
                            this.privateKey = new JcaPEMKeyConverter().getPrivateKey((PrivateKeyInfo) pemObj);
                        } else if (pemObj instanceof PrivateKey) {
                            this.privateKey = (PrivateKey) pemObj;
                        } else if (pemObj instanceof PublicKey) {
                            this.publicKey = (PublicKey) pemObj;
                        } else if (pemObj instanceof java.security.cert.Certificate) {
                            this.cert = (java.security.cert.Certificate) pemObj;
                            this.publicKey = cert.getPublicKey();
                        }
                    }
                }
                return;
                case "PKCS7": {
                    PKCS7 pkcs7 = new PKCS7(fis);
                    this.certChain = pkcs7.getCertificates();
                    this.cert = this.certChain[0];
                    this.publicKey = cert != null ? cert.getPublicKey() : null;
                }
                return;


                default: {
                    // pfx,p12 , jks 等
                    final KeyStore keyStore = KeyStore.getInstance(storeType);
                    keyStore.load(fis, storePass == null ? new char[0] : storePass.toCharArray());
                    this.privateKey = (PrivateKey) keyStore.getKey(keyAlias, keyPass == null ? null : keyPass.toCharArray());
                    this.cert = keyStore.getCertificate(keyAlias);
                    this.certChain = keyStore.getCertificateChain(keyAlias);
                    this.publicKey = cert != null ? cert.getPublicKey() : null;
                    if (publicKey == null) {
                    }
                }
            }
        } catch (Throwable ex) {
            throw new RuntimeException(ex);
        }
    }

    public java.security.cert.Certificate getCertificate() {
        return this.cert;
    }

    public java.security.cert.Certificate[] getCertificateChain() {
        return this.certChain;
    }

    private PublicKey _pubKeyFromPrivate;

    public PublicKey getPublicKey() {
        if (this.publicKey == null && this.privateKey != null) try {
            if (_pubKeyFromPrivate != null)
                return _pubKeyFromPrivate;
            return _pubKeyFromPrivate = generatePublicKey(this.privateKey);
        } catch (Throwable ex) {
            throw new RuntimeException(ex);
        }
        return this.publicKey;
    }

    public PrivateKey getPrivateKey() {
        return this.privateKey;
    }

    private PublicKey generatePublicKey(PrivateKey privKey) throws NoSuchAlgorithmException, InvalidKeySpecException {
        if (privKey instanceof RSAPrivateCrtKey) {
            BigInteger modulus = ((RSAPrivateCrtKey) privKey).getModulus();
            BigInteger publicExponent = ((RSAPrivateCrtKey) privKey).getPublicExponent();
            java.security.KeyFactory kf = java.security.KeyFactory.getInstance("RSA");
            RSAPublicKeySpec publicKeySpec = new RSAPublicKeySpec(modulus, publicExponent);
            return kf.generatePublic(publicKeySpec);
        }
        return null;
    }


    public KeyPair getKeyPair() {
        return new KeyPair(this.getPublicKey(), this.privateKey);
    }

	/*
	 *  "CN=caoyuwu,O=snsoft,O=snsoft,L=bj,ST=bj,C=CN",
        "CN=caoyuwu,O=snsoft,O=snsoft,L=bj,ST=bj,C=CN",
        50*365
        );
	 */

    static public X509Certificate loadX509Certificate(java.io.InputStream is) throws CertificateException {
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        return (X509Certificate) cf.generateCertificate(is);
    }

    public void saveKeyAsPEM(java.io.OutputStream out) {
        try (
                final JcaPEMWriter pemWriter = new JcaPEMWriter(new java.io.OutputStreamWriter(out));) {
            if (privateKey != null) {
                if (publicKey != null)
                    pemWriter.writeObject(new KeyPair(this.publicKey, this.privateKey));
                else
                    pemWriter.writeObject(this.privateKey);
            } else if (publicKey != null) {
                pemWriter.writeObject(this.publicKey);
            } else {
                throw new NullPointerException();
            }
        } catch (java.io.IOException ex) {
            throw new RuntimeException(ex);
        }
    }

    public void saveCertAsPEM(java.io.OutputStream out) {
        if (this.cert == null)
            throw new NullPointerException();
        try (final JcaPEMWriter pemWriter = new JcaPEMWriter(new java.io.OutputStreamWriter(out));) {
            pemWriter.writeObject(this.cert);
        } catch (java.io.IOException ex) {
            throw new RuntimeException(ex);
        }
    }

    public void save(java.io.OutputStream out,
                     java.io.InputStream in,
                     String storeType,
                     final String storePass,
                     final String keyAlias, final String keyPass) {
        try {
            KeyStore keyStore = KeyStore.getInstance(storeType);
            keyStore.load(in, storePass == null ? null : storePass.toCharArray());
            keyStore.setKeyEntry(keyAlias, privateKey, keyPass == null ? null : keyPass.toCharArray(),
                    cert == null ? null : new java.security.cert.Certificate[]{cert});
            keyStore.store(out, storePass == null ? null : storePass.toCharArray());
        } catch (Throwable ex) {
            throw new RuntimeException(ex);
        }
    }

}
